import os
import tensorflow as tf
from PIL import Image
import random

cwd = os.getcwd()
file_dir = 'F:\\001-python\\train_1\\'
filelist = []


def create_record_list():
    for file in os.listdir(file_dir):
        filelist.append(file)
        '''
        name = file.split(sep='.')
        lable_val = 0
        if name[0] == 'cat':
            lable_val = 0
        else:
            lable_val = 1
        img_path = file_dir + file
        img = Image.open(img_path)
        img = img.resize((208, 208))
        img_raw = img.tobytes()  # 将图片转化为原生bytes
        example = tf.train.Example(features=tf.train.Features(feature={
            "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[lable_val])),
            'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
        }))

        writer.write(example.SerializeToString())
        i=i+1
        print(i)
        '''


def create_record(filelist):
    random.shuffle(filelist)
    i = 0
    writer = tf.python_io.TFRecordWriter("catvsdogtrain.tfrecords")
    for file in filelist:
        name = file.split(sep='.')
        lable_val = 0
        if name[0] == 'cat':
            lable_val = 0
        else:
            lable_val = 1
        img_path = file_dir + file
        img = Image.open(img_path)
        img = img.resize((208, 208))
        img_raw = img.tobytes()  # 将图片转化为原生bytes
        example = tf.train.Example(features=tf.train.Features(feature={
            "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[lable_val])),
            'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
        }))

        writer.write(example.SerializeToString())

        print(name[1])

    writer.close()


if __name__ == '__main__':
    create_record_list()
    create_record(filelist)
